import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from scipy.fft import fft
from sklearn.feature_selection import SelectKBest, f_classif
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.svm import SVC
from imblearn.over_sampling import SMOTE
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.ensemble import RandomForestClassifier

# 1. 数据读取
train_data = pd.read_excel('exercise.xlsx')
test_data = pd.read_excel('test2.xlsx')
attachment_four = pd.read_excel('four.xlsx')

# 2. 数据预处理
train_data.columns = train_data.columns.astype(str)
test_data.columns = test_data.columns.astype(str)

X_train_raw = train_data.iloc[:, 4:]
X_test_raw = test_data.iloc[:, 4:]

scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train_raw)
X_test_scaled = scaler.transform(X_test_raw)

# 标签映射
waveform_label_map = {'正弦波': 1, '三角波': 2, '梯形波': 3}
y_train = train_data['waveform_label'].map(waveform_label_map)


# 3. 提取时域特征和频域特征
def extract_time_domain_features(X):
    features = pd.DataFrame()
    features['max'] = np.max(X, axis=1)
    features['min'] = np.min(X, axis=1)
    features['mean'] = np.mean(X, axis=1)
    features['std'] = np.std(X, axis=1)
    features['peak_to_peak'] = np.max(X, axis=1) - np.min(X, axis=1)
    return features

def extract_frequency_domain_features(X):
    features = pd.DataFrame()
    fft_features = fft(X, axis=1)
    features['fft_max_freq'] = np.max(np.abs(fft_features), axis=1)
    features['fft_energy'] = np.sum(np.abs(fft_features)**2, axis=1)
    return features

# 提取特征
train_time_features = extract_time_domain_features(X_train_scaled)
train_freq_features = extract_frequency_domain_features(X_train_scaled)
test_time_features = extract_time_domain_features(X_test_scaled)
test_freq_features = extract_frequency_domain_features(X_test_scaled)

X_train_features = pd.concat([train_time_features, train_freq_features], axis=1)
X_test_features = pd.concat([test_time_features, test_freq_features], axis=1)

# 4. 特征选择
selector = SelectKBest(f_classif, k=10)
X_train_selected = selector.fit_transform(X_train_features, y_train)
X_test_selected = selector.transform(X_test_features)

# 5. 使用 SMOTE 处理数据不平衡
smote = SMOTE(random_state=42)
X_train_resampled, y_train_resampled = smote.fit_resample(X_train_selected, y_train)

# 6. 使用网格搜索优化 SVM 模型
param_grid = {'C': [0.1, 1, 10, 100], 'gamma': ['scale', 'auto', 0.001, 0.01, 0.1, 1]}
grid_search = GridSearchCV(SVC(), param_grid, cv=5, scoring='accuracy')
grid_search.fit(X_train_resampled, y_train_resampled)
svm_model = grid_search.best_estimator_

# 7. 训练集和验证集划分，用于验证模型性能
X_train_split, X_val_split, y_train_split, y_val_split = train_test_split(X_train_selected, y_train, test_size=0.2, random_state=42)
y_val_pred = svm_model.predict(X_val_split)

# 打印分类报告和混淆矩阵
print("SVM 验证集性能评估")
print(classification_report(y_val_split, y_val_pred))
print(confusion_matrix(y_val_split, y_val_pred))

# 8. 预测测试集
y_test_pred = svm_model.predict(X_test_selected)

# 输出测试集预测结果
print(f"Test set length: {len(X_test_selected)}")
print(f"Prediction length: {len(y_test_pred)}")
print(f"Attachment four length: {len(attachment_four)}")

# 将预测结果写入附件四的文件中
if len(y_test_pred) != len(attachment_four):
    attachment_four.loc[:len(y_test_pred)-1, '附件二（80个样品）励磁波形分类结果'] = y_test_pred
else:
    attachment_four['附件二（80个样品）励磁波形分类结果'] = y_test_pred

# 导出结果至 Excel 文件
attachment_four.to_excel('four.xlsx', index=False)
print("结果已导出至 four.xlsx")

from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# 设置中文字体和负号显示
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体 (SimHei) 字体
plt.rcParams['axes.unicode_minus'] = False    # 解决坐标轴负号显示问题

# 绘制验证集的混淆矩阵
val_conf_matrix = confusion_matrix(y_val_split, y_val_pred)
fig, ax = plt.subplots(figsize=(8, 6))

# 使用 seaborn 绘制热力图以增强视觉效果
sns.heatmap(val_conf_matrix, annot=True, fmt='d', cmap='viridis', cbar=True,
            linewidths=1, linecolor='black', square=True, ax=ax)

# 设置标题和轴标签
plt.title("SVM 模型的混淆矩阵", fontsize=16)
plt.xlabel("预测标签", fontsize=14)
plt.ylabel("真实标签", fontsize=14)

# 设置轴刻度标签（根据实际类别名称替换标签）
labels = ['正弦波', '三角波', '梯形波']
ax.set_xticks(np.arange(len(labels)) + 0.5)
ax.set_yticks(np.arange(len(labels)) + 0.5)
ax.set_xticklabels(labels)
ax.set_yticklabels(labels)

# 调整热力图中格子大小和网格线
plt.gca().set_frame_on(False)
ax.set_xticks([x - 0.5 for x in range(1, len(labels))], minor=True)
ax.set_yticks([y - 0.5 for y in range(1, len(labels))], minor=True)
ax.grid(which='minor', color='black', linestyle='-', linewidth=2)  # 添加清晰的边界

# 显示图表
plt.show()




